#include "MLE.h"

MaximumLikelihoodEstimate::MaximumLikelihoodEstimate()
{
                                                      
}

MaximumLikelihoodEstimate::~MaximumLikelihoodEstimate()
{
     MLEAbund.clear();                                                      
}


int MaximumLikelihoodEstimate::check_assignment_accuracy(string tReadName, RDPIDQuery& RDPQuery, GBProcessing& GenBank, int TaxonomicLevel)
{
    size_t found=tReadName.find_last_of(".");
    string tTemplateName=RDPQuery.get_RDPID(tReadName.substr(found-8,8));
    if(AssignmentData[tReadName]!=GenBank.get_Taxonomy(tTemplateName, TaxonomicLevel))
        return 0;
    else
        return 1;  
}  

int MaximumLikelihoodEstimate::maximum_likelihood_estimate(map<string, map<string, double> >& ReadDistr, vector<string>& SelectedGroups)
{
    int TotalCount=0;
    int tSize=SelectedGroups.size();
    for(int i=0;i<tSize;i++)
        MLEAbund[SelectedGroups[i]]=0.0;
    map<string, map<string, double> >::iterator it=ReadDistr.begin();
    while(it!=ReadDistr.end())
    {
        TotalCount++;
        double tSum=0.0;
        for(int i=0;i<tSize;i++)
        {
            map<string, double>::iterator it2=(*it).second.find(SelectedGroups[i]);
            if(it2!=(*it).second.end())
                tSum+=(*it2).second;
        } //get Sum
        for(int i=0;i<tSize;i++)
        {
            map<string, double>::iterator it2=(*it).second.find(SelectedGroups[i]);
            if(it2!=(*it).second.end())
                MLEAbund[SelectedGroups[i]]+=(*it2).second/tSum;
        }
        it++;
    }
    /*
    //EM Part
    map<string, double> tAbund;
    double** T=new double*[TotalCount];
    for(int i=0; i<TotalCount; i++)
    {
        T[i]=new double[tSize];
        memset(T[i],0,sizeof(double)*tSize);
    }
    double Threshold=0.0; //for termination
    int tCount=0;
    int tIter=0;
    do
    {
        tCount=0;    
        it=ReadDistr.begin();
        while(it!=ReadDistr.end())
        {
            double tSum=0.0;
            for(int i=0;i<tSize;i++)
            {
                map<string, double>::iterator it2=(*it).second.find(SelectedGroups[i]);
                if(it2!=(*it).second.end())
                {
                    T[tCount][i]=(double)MLEAbund[SelectedGroups[i]]/TotalCount*(*it2).second;
                    tSum+=T[tCount][i];
                }
            } //get Sum
            for(int i=0;i<tSize;i++)
                T[tCount][i]/=tSum;
            tCount++;
            it++;            
        }  //E-step calculates Tij:=E(Zi=j|Data)
        tAbund=MLEAbund;
        for(int i=0;i<tSize;i++)
        {
            double tT=0.0;
            for(int j=0;j<TotalCount;j++)
                tT+=T[j][i];
            tT/=TotalCount;
            MLEAbund[SelectedGroups[i]]=tT;
            if(abs(MLEAbund[SelectedGroups[i]]-tAbund[SelectedGroups[i]])>Threshold)
                Threshold=abs(MLEAbund[SelectedGroups[i]]-tAbund[SelectedGroups[i]]);
        }
        tIter++;
    }while(Threshold>1e-8&&tIter<300);
    for(int i=0;i<TotalCount;i++)
        delete[] T[i];
    delete[] T;
    */
    return 1;
}
